package org.feng.config;

import dev.miku.r2dbc.mysql.MySqlConnectionConfiguration;
import dev.miku.r2dbc.mysql.MySqlConnectionFactory;
import lombok.extern.slf4j.Slf4j;
import org.springframework.context.annotation.Configuration;
import org.springframework.lang.NonNull;
import org.springframework.r2dbc.connection.lookup.AbstractRoutingConnectionFactory;
import reactor.core.publisher.Mono;

import javax.annotation.PostConstruct;
import javax.annotation.Resource;
import java.time.Duration;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Mysql 动态数据源
 *
 * @version v1.0
 * @author: fengjinsong
 * @date: 2022年06月14日 22时32分
 */
@Slf4j
@Configuration
public class MySqlDynamicDatasourceConfig extends AbstractRoutingConnectionFactory {
    /**
     * 默认租户ID
     */
    private static final String DEFAULT_MERCHANT_ID = "master";
    /**
     * 租户Key，切换数据源时，Context对象设置的key
     */
    public static final String MERCHANT_KEY = "merchantKey";
    /**
     * 数据源连接工厂映射
     */
    private static final Map<Object, Object> CONNECTION_FACTORY_CACHE_MAP = new HashMap<>(8);

    @Resource
    private MysqlDatasourceProperties mysqlDatasourceProperties;

    @Resource
    private MysqlDatasourcePropertiesV2 mysqlDatasourcePropertiesV2;

    @PostConstruct
    private void init() {
        final List<Property> propertyList = mysqlDatasourcePropertiesV2.getPROPERTY_LIST();
        for (Property property : propertyList) {
            MySqlConnectionFactory connectionFactory = MySqlConnectionFactory.from(MySqlConnectionConfiguration.builder()
                    .host(property.getHost())
                    .port(property.getPort())
                    .username(property.getUsername())
                    .password(property.getPassword())
                    .database(property.getDatabase())
                    .connectTimeout(Duration.ofSeconds(property.getConnectTimeout()))
                    .build());

            log.info("注册数据源 租户：{}  数据库名：{}", property.getMerchantKey(), property.getDatabase());
            CONNECTION_FACTORY_CACHE_MAP.put(property.getMerchantKey(), connectionFactory);
        }
        setTargetConnectionFactories(CONNECTION_FACTORY_CACHE_MAP);
        setDefaultTargetConnectionFactory(CONNECTION_FACTORY_CACHE_MAP.get(DEFAULT_MERCHANT_ID));
    }

    @NonNull
    @Override
    protected Mono<Object> determineCurrentLookupKey() {
        return Mono.deferContextual(Mono::just).handle((context, sink) -> {
            Object merchantKey = context.getOrDefault(MERCHANT_KEY, DEFAULT_MERCHANT_ID);
            log.info("使用数据源 {} HashCode {}", merchantKey, CONNECTION_FACTORY_CACHE_MAP.get(merchantKey).hashCode());
            assert merchantKey != null;
            sink.next(merchantKey);
        });
    }
}
